#include "importFunctions.hpp"
#include "gmshImport.hpp"
#include "meshTypeFunctions.hpp"
#include "rasterToField.hpp"
#include "vtkFunctions_import.hpp"
#include "loggerFunctions.hpp"

namespace hms
{


fs::path getVtkFilePath(
	bool resumeSimulation,
	const fs::path& caseDir,
	const std::string caseName
){
	using namespace vtk;
	/* if the simulation is to be resumed from an earlier state,
	 * set the vtk file path and start time accordingly */
	if (resumeSimulation){
		return lastWrittenFile( caseDir, caseName );
	} else {
		fs::path vtkFilePath { caseDir / caseName };
		tryMeshExtensions(vtkFilePath);
		return vtkFilePath;
	}
}

void setVtkFilePath(Settings& s){
	s.path(
		Settings::Path::vtkFile,
		getVtkFilePath( s.resumeSimulation, s.caseDir(), s.caseName() )
	);
}


void importAndOptimiseMesh(
	MeshVariant& meshVar,
	ImportSettings& importSettings
){
	std::shared_ptr<log::logger> logger { log::getLogger() };
	logger->debug("Importing mesh from\n"
		"{}\"{}\"", indent(), importSettings.mesh.file.string()
	);
	meshVar = gmsh::gmshToMesh( importSettings.mesh.file );
	importSettings.mesh.wasImported = true;

	if (importSettings.optimiseMeshType){
		logger->info("Attempting to optimise mesh type...");
		if ( optimiseMeshType(meshVar) ){
			std::visit( [&](const auto& mesh){
				using MeshType_new = remove_qualifiers<decltype(mesh)>;
				logger->info(
					"Mesh type was changed from {} to \"{}\".",
					bold("hms::Mesh"),
					bold("{}", type_name<MeshType_new>())
				);
			}, meshVar);
		} else {
			logger->info("Mesh already has optimal type.");
		}
	}
}


void createBackup( const fs::path& path, std::string_view message, bool dryRun ){
	fs::path backup { path };
	backup.replace_filename( path.filename().stem().string() + "_bak" );
	backup.replace_extension( path.extension() );
	log::info(
		"{}\nCreating backup file:"
		"\n{}\"{}\"",
		message, indent(), backup.string()
	);
	if (!dryRun)
		fs::copy( path, backup, fs::copy_options::update_existing );
}

void updateOrCreateDomainFile(
	Settings& s,
	xml::XMLDocument* domainDoc,
	const MeshVariant& meshVar
){
	using namespace vtk;
	bool domainDataFound {false};
	bool vtkFileFound { fs::is_regular_file( s.vtkFilePath() ) };
	if ( vtkFileFound ){
		/* read existing comain file in */
		readXML( domainDoc, s.vtkFilePath() );
		/* create backup file before changing it */
		createBackup( s.vtkFilePath(),
			"Existing vtk file found. Mesh portion is about to be overwritten.",
			s.dryRun
		);
		/* write imported mesh to new xml and copy domain fields over */
		std::unique_ptr<xml::XMLDocument> newDomainDoc { meshToVtk(meshVar) };
		/* this may throw if the user's xml doc is broken. In that case,
		 * we can simply create a default file, so we catch and discard
		 * that error. */
		auto getPiece = [&]( xml::XMLDocument* doc ) -> xml::XMLElement* {
			xml::XMLElement* vtk { getVtkNode(doc) };
			/* the child node's name is the value of the attribute "type" */
			const auto& vtkType { vtk->Attribute("type") };
			return getChildNodeByName(
				getChildNodeByName(vtk, vtkType), "Piece"
			);
		};
		try{
			const xml::XMLElement* cellData {
				getPiece(domainDoc)->FirstChildElement("CellData")
			};
			if (cellData){
				xml::XMLElement* newPiece { getPiece(newDomainDoc.get()) };
				xml::XMLNode* cellDataCopy {
					cellData->DeepClone( newDomainDoc.get() )
				};
				newPiece->InsertEndChild(cellDataCopy);
				/* delete the old vtk file, because of the possibly
				 * differing file ending, which may cause the old file to
				 * be imported on the next run */
				if (!s.dryRun){
					fs::remove( s.vtkFilePath() );
					replaceExtension( s.vtkFilePath(), meshVar );
					toFile( newDomainDoc.get(), s.vtkFilePath() );
				}
				domainDataFound = true;
			}
		} catch ( const ParseError& e ){
			std::cout << "Ignore this error:\n\t" << e.what() << "\n";
		}
	}
	if ( !vtkFileFound || !domainDataFound ){
		replaceExtension(s.vtkFilePath(), meshVar);

		log::info(
			"No vtk file found, or no domain field data found in vtk file.\n"
			"Creating default vtk file:"
			"\n{}\"{}\"",
			indent(), s.vtkFilePath().string()
		);
		meshToVtk(domainDoc, meshVar);
		std::visit( [&](const auto& mesh){
			using MeshType = remove_qualifiers<decltype(mesh)>;
			writeDefaultDomainFields<SWESolver<MeshType>>(domainDoc);
		}, meshVar);
		if (!s.dryRun)
			toFile(domainDoc, s.vtkFilePath(), meshVar);

		s.dummyFileCreated = true;
		s.dummyVtkFile = true;
	}
}


void updateOrCreateBoundaryFile(
	Settings& s,
	xml::XMLDocument* boundaryDoc,
	const MeshVariant& meshVar
){
	const fs::path& boundaryPath { s.path(Settings::Path::boundaryFile) };
	using namespace vtk;
	if ( fs::is_regular_file(boundaryPath) ){
		/* create backup file before changing it */
		createBackup( boundaryPath,
			"Existing boundary file found.\n"
			"Patch vertices may be overwritten and extraneous patches are "
			"removed, but boundary conditions for existing patches are "
			"copied to new file.",
			s.dryRun
		);
		readXML(boundaryDoc, boundaryPath);
		boundaryToXML(boundaryDoc, meshVar);
		/* the boundary file might contain boundary patches which the mesh
		 * doesn't contain. But due to potentially new vertex indices,
		 * these are not necessarily valid. */
		removeNonexistentPatches(boundaryDoc, meshVar);
		if (!s.dryRun)
			toFile(boundaryDoc, boundaryPath);
	} else {
		log::info(
			"No boundary file found.\n"
			"Boundary patches may have been defined "
			"during mesh import/generation,\nbut {}.\n"
			"Writing template to file:"
			"\n{}\"{}\"",
			bold( "boundary conditions must be defined in boundary file" ),
			indent(), boundaryPath.string()
		);
		boundaryToXML(boundaryDoc, meshVar);
		writeDefaultBoundaryConditions(boundaryDoc, meshVar);
		if (!s.dryRun)
			toFile(boundaryDoc, boundaryPath);

		s.dummyFileCreated = true;
	}
}


MeshVariant readMesh(
	Settings& s,
	xml::XMLDocument* domainDoc,
	xml::XMLDocument* boundaryDoc
){
	using namespace vtk;
	// bool vtkFileFound { fs::is_regular_file( s.vtkFilePath() ) };
	log::debug(
		"Reading mesh from files:"
		"\n{0}\"{1}\""
		"\n{0}\"{2}\"",
		indent(),
		s.vtkFilePath().string(),
		s.boundaryPath().string()
	);

	readXML( domainDoc, s.vtkFilePath() );
	getBoundaryDoc(boundaryDoc, s.caseDir(), s.dummyFileCreated, s.dryRun);
	MeshVariant meshVar { extensionToMeshType( s.vtkFilePath() ) };

	toMesh(domainDoc, boundaryDoc, meshVar);
	return meshVar;
}

MeshVariant readMesh( Settings& s ){
	xml::XMLDocument domainDoc, boundaryDoc;
	return readMesh(s, &domainDoc, &boundaryDoc);
}

void writeMesh( Settings& s, const MeshVariant& meshVar ){
	xml::XMLDocument domainDoc, boundaryDoc;
	/* check whether domain and boundary files are available, otherwise write
	 * default files as templates.
	 * If mesh was neither imported nor generated,
	 * it was only read and therefore does not need to be written.
	 */
	if ( s.import().mesh.wasImported || s.import().generateMesh ){
		s.logger().debug("Writing mesh to domain and boundary file...");
		/* domain file */
		updateOrCreateDomainFile( s, &domainDoc, meshVar );
		/* boundary */
		updateOrCreateBoundaryFile( s, &boundaryDoc, meshVar );
	}
}

template<typename FormatType>
void noImportMsg(
	const Settings& settings,
	typename ImportSettings::ImportItem<FormatType>& iniSetting,
	typename Settings::ImportItem cliOption,
	std::string_view what
){
	auto level { log::Level::off };
	if ( iniSetting.import != ImportSettings::Flag::never )
		level = log::Level::debug;
	if ( settings.importItem.value() == cliOption )
		level = log::Level::warn;

	if ( level != log::Level::off )
		settings.logger().log(level,
			"{} was not imported (import flag: \"{}\")",
			what,
			magic_enum::enum_name( iniSetting.import )
		);
}

MeshVariant getMesh( Settings& s ){
	using namespace vtk;

	bool vtkFileFound { fs::is_regular_file( s.vtkFilePath() ) };
	/* import files if required */
	bool importMesh { decideImport( s.import().mesh, s.vtkFilePath() ) };
	if (!importMesh)
		noImportMsg(s, s.import().mesh, Settings::ImportItem::mesh, "Mesh");

	/* or generate mesh */
	const bool& generateMesh { s.import().generateMesh };
	/* check for conflicting options */
	if ( (importMesh || generateMesh) && s.resumeSimulation )
		throw InputError( fmt::format(
			"{}. Cannot import/generate mesh with setting [time] startAt=last, "
			"due to likely mesh/data mismatch.",
			bold(red("Conflicting settings"))
		) );

	/* now generate, import, or read mesh from the vtk and boundary files */
	MeshVariant meshVar;

	if ( importMesh ){
		s.logger().info("Importing mesh...");
		importAndOptimiseMesh( meshVar, s.import() );
	} else if ( generateMesh ){
		s.logger().info("Generating mesh...");
		s.meshGen( getMeshGenSettings(s.iniReader) );
		meshVar = hms::generateMesh( s.meshGen() );
	} else if (vtkFileFound) {
		s.logger().info("Reading mesh...");
		meshVar = readMesh(s);
	} else {
		throw InputError(
			"No mesh generation/import requested and no vtk file found!"
		);
	}
	std::visit([](const auto& mesh){ mesh.printInfo(); }, meshVar);

	return meshVar;
}



ArrayXXs importRaster( const RasterItem& item, const MeshVariant& meshVar ){
	ArrayXXs fieldValues;
	using namespace hms::raster;
	Raster raster { readRaster(item.file, item.format, item.formatFile) };
	std::visit( [&](const auto& mesh){
		fieldValues = rasterToFieldValues(raster, mesh);
	}, meshVar);
	return fieldValues;
}



std::optional<ArrayXXs> readDEM(
	Settings& settings,
	const MeshVariant& meshVar
){
	const RasterItem& dem { settings.import().dem };
	bool importZField { decideImport(
		dem, settings.vtkFilePath(), settings.dummyVtkFile
	) };

	if ( importZField ){
		log::debug(
			"Importing DEM from:\n{}\"{}\"", indent(), dem.file.string()
		);
		ArrayXXs zValues { importRaster(dem, meshVar) };

		settings.import().dem.wasImported = true;

		return zValues;
	}
	noImportMsg(
		settings, settings.import().dem, Settings::ImportItem::dem, "DEM"
	);

	return std::nullopt;
}

void writeDEM(
	ArrayXXs&& zValues,
	Settings& settings,
	const MeshVariant& meshVar,
	xml::XMLDocument* domainDoc,
	xml::XMLDocument* boundaryDoc
){
	using namespace vtk;
	std::visit( [&](const auto& mesh){
		using MeshType = remove_qualifiers<decltype(mesh)>;

		SWESolver solver { mesh, Time{} };
		Field<MeshType>& z { solver.zField() };
		z.values() = std::move( zValues );
		/* set the boundary conditions of the z field to staticValue,
			* because they do not need to be changed anymore */
		for ( const BoundaryPatch& patch : mesh.boundaryPatches() ){
			z.setBoundaryCondition( StaticValue{z, patch, false} );
		}
		writeDomainField( domainDoc, solver, solver.zFieldName() );
		writeBoundaryConditions( boundaryDoc, solver, solver.zFieldName() );

	}, meshVar );

	if (!settings.dryRun){
		toFile( domainDoc, settings.vtkFilePath() );
		toFile( boundaryDoc, settings.boundaryPath());
	}
}

void importDEM(
	Settings& settings,
	const MeshVariant& meshVar,
	xml::XMLDocument* domainDoc,
	xml::XMLDocument* boundaryDoc
){
	std::optional<ArrayXXs> zValues { readDEM(settings, meshVar) };
	if (zValues)
		writeDEM(
			std::move(*zValues), settings, meshVar, domainDoc, boundaryDoc
		);
}


template<typename FieldName>
void deleteFieldNode(
	xml::XMLNode* root,
	const char* childName,
	FieldName fieldName
){
	using namespace vtk;
	xml::XMLElement* fieldNode {
		vtk::internal::getChildNodeByAttributeValue(
			root, childName,
			"Name", magic_enum::enum_name(fieldName)
	) };
	if ( fieldNode ) root->DeleteChild(fieldNode);
}


std::optional<ArrayXXs> readRoughness(
	Settings& settings,
	const MeshVariant& meshVar
){
	// using namespace vtk;
	RasterItem& roughness { settings.import().roughness };
	auto cfp { Settings::Path::constFields };
	settings.path( cfp, vtk::getConstantFieldsPath( settings.caseDir() ) );
	const fs::path& constFieldsPath { settings.path(cfp) };

	std::shared_ptr<spdlog::logger> logger { log::getLogger() };
	logger->debug( "Constant fields file:\n{}{}",
		indent(),
		constFieldsPath.string()
	);

	if ( decideImport(roughness, constFieldsPath) ){
		if ( !settings.swe().useFriction ){
			logger->warn(
				"Settings mismatch:\n"
				"{}Roughness data import requested, but friction is disabled.",
				indent()
			);
		}
		logger->debug(
			"Importing roughness data from:\n{}\"{}\"",
			indent(), roughness.file.string()
		);
		/* read raster data and scale it to field values */
		ArrayXXs roughnessValues { importRaster(roughness, meshVar) };

		roughness.wasImported = true;

		return roughnessValues;
	}
	noImportMsg(
		settings,
		settings.import().roughness,
		Settings::ImportItem::roughness, "Roughness data"
	);

	return std::nullopt;
}

void writeRoughness(
	ArrayXXs&& fieldVals,
	Settings& settings,
	const MeshVariant& meshVar
){
	using namespace vtk;
	xml::XMLDocument constDoc;
	getConstantFieldsDoc(
		&constDoc, settings.caseDir(),
		settings.dummyFileCreated, settings.dryRun
	);

	std::visit( [&](const auto& mesh){
		SWESolver solver { mesh, Time{} };
		using MeshType = remove_qualifiers<decltype(mesh)>;
		auto fieldName = SWESolver<MeshType>::ConstantFieldName::ManningsN;

		solver.n() = ConstantField<NonUniform, MeshType>{
			mesh, std::move( fieldVals )
		};

		xml::XMLElement* root { getRootOrInsert(&constDoc, constNode) };
		deleteFieldNode(root, constChild, fieldName);

		writeConstantField( root, solver, fieldName);
	}, meshVar );

	if (!settings.dryRun)
		toFile( &constDoc, settings.constFieldsPath() );
}

void importRoughness(
	Settings& settings,
	const MeshVariant& meshVar
){
	std::optional<ArrayXXs> roughnessVals { readRoughness(settings, meshVar) };
	if (roughnessVals)
		writeRoughness( std::move(*roughnessVals), settings, meshVar);
}



std::optional<TimeSeries> readRain(
	Settings& settings,
	[[maybe_unused]] const MeshVariant& meshVar
){
	using namespace vtk;
	const RasterItem& rain { settings.import().rain };
	std::shared_ptr<spdlog::logger> logger { log::getLogger() };

	auto msg = [&](){
		noImportMsg(
			settings, settings.import().rain,
			Settings::ImportItem::rainfall, "Rainfall data"
		);
	};

	/* if no import is supposed to happen, return early */
	if ( rain.import == ImportSettings::Flag::never ){
		msg();
		return std::nullopt;
	}

	auto sfp { Settings::Path::sourceFields };

	settings.path( sfp, getSourceFieldsPath( settings.caseDir() ) );

	const fs::path& sourceFieldsPath { settings.path(sfp) };
	logger->debug(
		"Source fields file:\n{}{}",
		indent(), sourceFieldsPath.string()
	);

	std::string rainFilename {
		rain.file.filename().replace_extension("").string()
	};

	/* compare file names. If the file name starts with what's specified 
	 * in the settings, then it's a match.
	 * Then it should continue with (an underscore and) a number, which 
	 * is the time value in seconds. */
	auto hasTime = [&]( const fs::path& p ) -> std::optional<scalar> {
		std::regex pattern {
			rainFilename + "_?([0-9]+)\\" + rain.file.extension().string()
		};
		std::smatch match_out;
		std::string filename { p.filename() };
		if ( std::regex_match(
			filename, match_out, pattern
		) ){
			/* first entry is whole string, second is piece match */
			std::ssub_match timeString { match_out[1] };
			std::stringstream conv { timeString.str() };
			scalar t;
			if ( conv >> t )
				return t;

		}
		return std::nullopt;
	};

	/* if any of the files match the flag specified in the settings,
	 * the imported time series is written to the source fields file */
	bool writeData {false};

	/* copy the raster item, so that we can change the path to point to a file
	 * with a time value */
	RasterItem timeItem {rain};

	TimeSeries timeSeries;

	fmt::memory_buffer out;

	/* iterate over all files in the import directory and compare them to the
	 * specified file name */
	for( auto&& filepath : fs::directory_iterator( rain.file.parent_path() ) ){
		if ( auto timeOpt = hasTime(filepath) ){
			timeItem.file = std::move(filepath);

			/* decideImport is only called until it returns true once. */
			if ( !writeData && decideImport(timeItem, sourceFieldsPath) ){
				writeData = true;
			}

			timeSeries.emplace_back(
				timeOpt.value(), importRaster(timeItem, meshVar)
			);
			fmt::format_to(std::back_inserter(out),
				"\n{}\"{}\"", indent(), timeItem.file.string()
			);
		}
	}

	std::sort(timeSeries.begin(), timeSeries.end(),
		[](const auto& pair1, const auto& pair2){
			return pair1.first < pair2.first;
		} );

	/* write data to source fields file, from where it can natively be imported */
	if ( writeData ){
		logger->debug( "Importing rainfall data from:{}", fmt::to_string(out) );
		using SF = typename SWESettings::SourceFieldName;
		if ( !settings.swe().useSourceFields[static_cast<int>(SF::Rainfall)] ){
			logger->warn(
				"Settings mismatch:\n"
				"{}Rainfall data import requested, "
				"but rainfall source term is disabled.",
				indent()
			);
		}
		settings.import().rain.wasImported = true;

		return timeSeries;
	}
	msg();

	return std::nullopt;
}

void writeRain(
	TimeSeries&& rainfall,
	Settings& settings,
	const MeshVariant& meshVar
){
	using namespace vtk;
	xml::XMLDocument sourceDoc;
	getSourceFieldsDoc(
		&sourceDoc, settings.caseDir(),
		settings.dummyFileCreated, settings.dryRun
	);

	std::visit( [&](const auto& mesh){
		using MeshType = remove_qualifiers<decltype(mesh)>;
		auto fieldName = SWESolver<MeshType>::SourceFieldName::Rainfall;

		SWESolver solver { mesh, Time{} };
		SourceField<NonUniform, MeshType> rainField;
		rainField.mesh( mesh );
		rainField.timeSeries( std::move( rainfall ) );
		rainField.useInterpolation( settings.import().rain.interpolate );

		solver.r() = std::move(rainField);
		solver.settings().useSourceFields[static_cast<int>(fieldName)] = true;

		xml::XMLElement* root { getRootOrInsert(&sourceDoc, sourceNode) };
		deleteFieldNode(root, sourceChild, fieldName);
		writeSourceField( root, solver, fieldName);
	}, meshVar );

	if (!settings.dryRun)
		toFile( &sourceDoc, settings.sourceFieldsPath() );
}

void importRain(
	Settings& settings,
	const MeshVariant& meshVar
){
	std::optional<TimeSeries> timeVals { readRain(settings, meshVar) };
	if ( timeVals )
		writeRain( std::move(*timeVals), settings, meshVar );
}

void importData(Settings& settings){
	using namespace vtk;

	log::printSection(settings.logger(), "Data import", log::Level::info);

	if ( settings.runBenchmarkCase ){
		throw InputError(
			"Flag for running benchmark is set. "
			"Benchmarks are run on fixed settings, not on imported data."
		);
	}

	MeshVariant meshVar;
	xml::XMLDocument domainDoc, boundaryDoc;

	/* get settings for file import (mesh, raster data) */
	setImportSettings(settings);

	using SI = Settings::ImportItem;
	SI what { settings.importItem.value() };

	/* If mesh import/generation is performed, the results are written out.
	 * If not, it is expected that a VTK file already exists */
	if (
		   what == SI::all
		|| what == SI::mesh
		|| !fs::is_regular_file( settings.vtkFilePath() )
	){
		meshVar = getMesh(settings);
		writeMesh(settings, meshVar);
		readXML( &domainDoc, settings.vtkFilePath() );
		readXML( &boundaryDoc, settings.boundaryPath() );
	} else {
		/* In any case, both the mesh variant and its vtk representation
		 * are required for importing. */
		meshVar = readMesh(settings, &domainDoc, &boundaryDoc);
	}
	if ( what == SI::all || what == SI::dem ){
		importDEM(settings, meshVar, &domainDoc, &boundaryDoc);
	}
	if ( what == SI::all || what == SI::rainfall ){
		importRain(settings, meshVar);
	}
	if ( what == SI::all || what == SI::roughness ){
		importRoughness(settings, meshVar);
	}
	settings.logger().info("Import completed, program will now exit.");
}


} // namespace hms
